from app.models.base import BaseModel, db
from enum import Enum
from datetime import datetime
from typing import Dict, Any

class TaskStatus(Enum):
    PENDING = 'pending'
    RUNNING = 'running'
    SUCCESS = 'success'
    FAILURE = 'failure'

class TaskType(Enum):
    DOCUMENT_PROCESS = 'document_process'
    EMBEDDING_GENERATE = 'embedding_generate'

class Task(BaseModel):
    """任务记录"""
    __tablename__ = 'task'

    id = db.Column(db.Integer, primary_key=True)
    celery_id = db.Column(db.String(36), unique=True)
    type = db.Column(db.String(50), nullable=False)
    status = db.Column(db.String(20), default=TaskStatus.PENDING.value)
    progress = db.Column(db.Float, default=0.0)  # 0-100
    result = db.Column(db.JSON)
    error = db.Column(db.Text)
    
    def to_dict(self) -> Dict[str, Any]:
        result = super().to_dict()
        result['type'] = self.type
        result['status'] = self.status
        return result

    @classmethod
    def create_task(
        cls,
        celery_id: str,
        task_type: TaskType
    ) -> 'Task':
        task = cls(
            celery_id=celery_id,
            type=task_type.value
        )
        db.session.add(task)
        db.session.commit()
        return task

    def update_status(
        self,
        status: TaskStatus,
        progress: float = None,
        result: Dict = None,
        error: str = None
    ) -> None:
        self.status = status.value
        if progress is not None:
            self.progress = progress
        if result is not None:
            self.result = result
        if error is not None:
            self.error = error
        db.session.commit() 